import traceback
import collections
import sys

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torcheeg.transforms import MeanStdNormalize
from models import *
from src.preprocess_v2 import read_edf
from src.tags import ToTags

MODELPATH = '../model/model_MyLSTM_v2.1.pt'

def load_model(model_path:str, device:str='cpu') -> nn.Module:
    '''加载模型'''
    model = MyLSTM()
    model.to(device)
    # hook_model(model)
    src_state_dict:dict[str, torch.Tensor] = torch.load(model_path, map_location=device)
    state_dict = collections.OrderedDict()
    for key, value in src_state_dict['state_dict'].items():
        if key.startswith('model.'):
            state_dict[key.replace('model.', '')] = value
    model.load_state_dict(state_dict)
    return model

def load_data(edf_path:str, window:int=100, step:int=50) -> Tuple[Iterator[np.ndarray], int]:
    '''加载数据'''
    raw, SRR = read_edf(edf_path, new_freq=100)
    ori_data = raw.get_data()[:16]
    n_times = raw.n_times
    raw.close()

    transform = MeanStdNormalize(axis=1)
    data_gen = (transform(eeg=ori_data[:, i:i+window])['eeg'] \
                for i in range(0, n_times-window+1, step))
    start_gen = (i for i in range(0, n_times-window+1, step))
    return data_gen, start_gen, SRR

def pred(datas, starts, model:nn.Module, SRR:int, out_path:str, device:str='cpu', 
         verbose:bool=True) -> None:
    '''预测'''
    with torch.no_grad():
        model.eval()
        pred_loader = DataLoader(list(datas), batch_size=64)
        y_pred = np.array([], dtype=int)
        for data in pred_loader:
            data: torch.Tensor
            batch_yp:torch.Tensor = model(data.to(device).float())
            batch_yp = batch_yp.max(dim=1)[1]
            y_pred = np.append(y_pred, batch_yp.numpy())

    if verbose:
        print(np.unique(y_pred, return_counts=True))
    tt = ToTags(100, label_map={1: '尖慢波'})
    tt.set_SRR(SRR)
    tt.collect_event(y_pred, starts, converted=False)
    tt.to_tags(out_path=out_path)

def main():
    edf_path = sys.argv[1]
    model = load_model(MODELPATH)
    data_gen, start_gen, SRR = load_data(edf_path)
    out_path = edf_path.lower().replace('.edf', '')+'.tags'
    pred(data_gen, start_gen, model, SRR, out_path)

if __name__=='__main__':
    try:
        main()
        print("eegpysuccess")
    except Exception as e:
        traceback.print_exc()
        print('eegpyerror')
